from hmmlearn import hmm
import numpy as np

transition_probability = np.array([
  [0.5, 0.2, 0.3],
  [0.3, 0.5, 0.2],
  [0.2, 0.3, 0.5]
])

emission_probability = np.array([
  [0.5, 0.5],
  [0.4, 0.6],
  [0.7, 0.3]
])

start_probability = np.array([0.2, 0.3, 0.5])
seen = np.array([[0, 1, 0, 1]]).T
model = hmm.MultinomialHMM(n_components=3)
model.startprob_=start_probability
model.transmat_=transition_probability
model.emissionprob_=emission_probability
logprob, box = model.decode(seen, algorithm="viterbi")
print(logprob, box)




def viterbi(A, B, pi, O):
  '''
  :param A: 状态转移概率矩阵：一个状态转移到另一个状态的概率
  :param B: 观测概率矩阵：行表示某个状态，列表示某个观测，Bij表示在这个状态下产生这个观测的概率,
  设第一列观测序号为0，第二列观测的序号为1，...
  :param pi: 初始状态概率向量：以某个状态开始的概率,pi的长度即为可选的状态的个数
  :param O: 观测序列，由0，1，2...组成，O的长度即为观测次数
  :return: 最优路径I:一个数组，表示每个时刻最可能选择的状态的编号i，为0,1,2...
  '''

  # 第一步：初始化:变量
  # k表示观测次数
  k = O.shape[0]
  # n表示可选状态的数目
  n = pi.shape[0]
  # node记录时刻t时，状态为i的所有路径中概率最大的路径的上一个结点，即为上一个状态的编号0，1，2...，大小为 k * n,每个元素初始化为0
  node = [[0 for i in range(n)] for j in range(k)]
  # delta记录时刻t时，状态i观测到O[t]的概率，大小为k * n，每个元素初始化为0
  delta = [[0 for i in range(n)] for j in range(k)]

  # 第二部：单独初始化时刻1，时刻1的观测为O[0]，
  for i in range(n):
    delta[0][i] = pi[i] * B[i][O[0]]

  # 第三步：递推计算:时刻2，3，4，....
  for t in range(1, k):
    for i in range(n):
      # 计算t时刻，状态i，在上一个时刻各个状态下得到的观测的基础上，该时刻i状态得到新观测值的最大概率
      # 即计算delta[t][i]
      # 计算node[t][i]，记录的是前一个状态j
      prob_max = 0
      j_max = 0
      for j in range(n):
        p = delta[t - 1][j] * A[j][i]
        if p > prob_max:
          prob_max = p
          j_max = j
      delta[t][i] = prob_max * B[i][O[t]]
      node[t][i] = j_max

  # 第四步：找到最优路径的终点
  # I为返回结果，是最优路径序列
  I = []
  i = np.argmax(delta[k - 1])
  I.append(i)

  # 第五步：由最优路径的终点回溯找到其他结点
  for t in range(k - 2, -1, -1):
    I.insert(0, node[t + 1][i])
    i = node[t + 1][i]

  return np.array(I)

res = viterbi(transition_probability, emission_probability, start_probability, seen)
print(res)